import divisionBlock from './gpuAccelerate/divisionBlock'
import tf from './tf'

interface CreateImageBlocksOptions {
  width: number
  height: number
  data: Float32Array //数据
  blockSize: number //一个块的尺寸 （宽度 ，高度）
  channel: number //通道数
  /**
   * padding，每个块四周的padding
   */
  padding: number
}
interface ImageBlockOptions {
  width: number
  height: number
  data: Float32Array //数据
  channel: number //通道数
}
/**
 * 获取最近边缘位置
 * @param pos
 * @param width
 * @param height
 * @returns
 */
function getLatelyEdgePos(
  pos: { x: number; y: number },
  width: number,
  height: number
) {
  if (pos.x < 0) {
    pos.x = 0
  }
  if (pos.y < 0) {
    pos.y = 0
  }
  if (pos.x >= width) {
    pos.x = width - 1
  }
  if (pos.y >= height) {
    pos.y = height - 1
  }
  return { ...pos }
}
/**
 * 图片块，将可以将一张图片切割成多个块，每个块可以单独处理
 * 处理完后再进行合成
 */
export default class ImageBlocks {
  width: number
  height: number
  list: ImageBlock[]
  blockRows: number
  blockCols: number
  channel: number
  padding: number
  data: Float32Array
  blockSize: number
  private onGpuDestroy?: () => void
  constructor(options: CreateImageBlocksOptions) {
    //初始化宽度
    this.width = options.width
    this.height = options.height
    this.list = []
    this.blockRows = Math.ceil(this.height / options.blockSize)
    this.blockCols = Math.ceil(this.width / options.blockSize)
    this.channel = options.channel
    this.padding = options.padding
    this.data = options.data
    this.blockSize = options.blockSize
  }
  async init() {
    //GPU初始化
    const { value, destroy } = await divisionBlock(
      this.data,
      this.height,
      this.width,
      this.channel,
      this.blockSize,
      this.padding
    )
    if (value) {
      //支持webGPU并且初始化成功
      // console.log('支持webGPU并且初始化成功')
      for (let i = 0; i < this.blockRows; i++) {
        for (let j = 0; j < this.blockCols; j++) {
          const index = i * this.blockCols + j
          this.list.push(
            new ImageBlock({
              width: this.blockSize + this.padding * 2,
              height: this.blockSize + this.padding * 2,
              data: value[index],
              channel: this.channel,
            })
          )
        }
      }
      this.onGpuDestroy = destroy
      // return
    }

    // CPU初始化
    const getPosData = (x: number, y: number, deep = false) => {
      const pos = (y * this.width + x) * this.channel
      //超出范围
      if (pos + this.channel > this.data.length || pos < 0) {
        //如果是4通道
        // if (options.channel == 4) {
        //   return [0, 0, 0, 1]
        // }
        if (deep) {
          debugger
        }
        const point = getLatelyEdgePos({ x, y }, this.width, this.height)
        return getPosData(point.x, point.y, true)
      }
      return this.data.slice(pos, pos + this.channel)
    }
    this.list = []
    for (let i = 0; i < this.blockRows; i++) {
      for (let j = 0; j < this.blockCols; j++) {
        const tsY = i * this.blockSize - this.padding
        const tsX = j * this.blockSize - this.padding
        let blockBuffer: number[] = []
        let b = 0
        for (let y = tsY; y < tsY + this.blockSize + 2 * this.padding; y++) {
          //遍历列
          for (let x = tsX; x < tsX + this.blockSize + 2 * this.padding; x++) {
            const data = getPosData(x, y)
            blockBuffer.push(...data)
            b++
          }
        }
        this.list.push(
          new ImageBlock({
            width: this.blockSize + this.padding * 2,
            height: this.blockSize + this.padding * 2,
            data: new Float32Array(blockBuffer),
            channel: this.channel,
          })
        )
      }
    }
  }

  destroy() {
    this.onGpuDestroy && this.onGpuDestroy()
  }

  /**
   *
   * @param list 图片块列表 (块内宽高包含了padding)
   * @param width 整个图片的宽度
   * @param height 整个图片的高度
   * @param padding 块的padding
   * @returns
   */
  static listToTensor(
    list: ImageBlock[],
    width: number,
    height: number,
    padding: number
  ): tf.Tensor3D {
    if (list.length == 0) throw new Error('list is empty')
    // 创建一个Uint8ClampedArray，长度为width * height * channel
    const data = new Float32Array(width * height * 3)
    const bWidth = list[0].width
    const bHeight = list[0].height
    const blockCols = Math.ceil(width / (bWidth - padding * 2))
    // const blockRows = Math.ceil(height / blockHeight)
    const setPosData = (
      x: number,
      y: number,
      value: number[],
      distance: number
    ) => {
      // 计算位置
      const pos = (y * width + x) * 3
      //超出范围
      if (pos > data.length) return
      // 赋值
      for (let i = 0; i < 3; i++) {
        //边缘覆盖
        if (data[pos + i] != 0 && distance < padding && padding != 0) {
          const proportion = distance / padding
          const tmp = value[i] * proportion
          data[pos + i] = tmp + data[pos + i] * (1 - proportion)
        } else {
          data[pos + i] = value[i]
        }
      }
    }
    // 定义一个函数，用于获取点到块边缘的距离
    const getPointToBlockEdgeDistance = (
      x: number,
      y: number,
      width: number,
      height: number
    ) => Math.min(x, width - x, y, height - y)
    for (let i = 0; i < list.length; i++) {
      const block = list[i]
      const tsY = Math.floor(i / blockCols) * (bHeight - padding * 2)
      const tsX = (i % blockCols) * (bWidth - padding * 2)
      for (let y = 0; y < bHeight; y++) {
        for (let x = 0; x < bWidth; x++) {
          const blockPos = { x: tsX - padding, y: tsY - padding }
          const distance = getPointToBlockEdgeDistance(x, y, bWidth, bHeight)
          const value = block.getData(x, y)
          setPosData(x + blockPos.x, y + blockPos.y, value, distance)
        }
      }
    }
    return tf.tensor3d(data, [height, width, 3], 'float32')
  }
}

export class ImageBlock {
  width: number
  height: number
  data: Float32Array
  channel: number //通道数
  constructor(options: ImageBlockOptions) {
    this.width = options.width
    this.height = options.height
    this.data = options.data
    this.channel = options.channel
  }

  /**
   * 将一个tensor转成ImageBlock
   * @param tensor [? batch=1,height,width,channel]
   * @returns
   */
  static async fromTensor(tensor: tf.Tensor3D | tf.Tensor4D) {
    const shape = [...tensor.shape]
    if (tensor.dtype == 'int32') {
      debugger
      tensor = tf.tidy(() => tensor.cast('float32').div(255))
    }
    if (shape.length == 4) {
      if (shape[0] != 1) throw new Error('shape[0] batch must be 1')
      shape.shift()
    }
    // if (tensor.dtype === 'float32') {
    //   //从float32转uint8
    //   //@ts-ignore
    //   tensor = tf.tidy(() => {
    //     return tf.clipByValue(tensor.mul(255).round().cast('int32'), 0, 255)
    //   })
    // }
    return new ImageBlock({
      width: shape[1],
      height: shape[0],
      data: (await tensor.data()) as Float32Array,
      channel: shape[2],
    })
  }

  public toInputTensor(): tf.Tensor3D {
    return tf.tidy(() => {
      return tf.tensor3d(
        this.data,
        [this.height, this.width, this.channel],
        'float32'
      )
    })
  }

  public getData(x, y) {
    const pos = (y * this.width + x) * this.channel
    if (pos > this.data.length) {
      console.warn('pos out of range')
      if (this.channel == 4) {
        return [0, 0, 0, 255]
      }
      return new Array(this.channel).fill(0)
    }
    return [...this.data.slice(pos, pos + this.channel)]
  }
}

// function toUint8ClampedArray(buffer: {
//   [index: number]: number
//   length: number
// }) {
//   const data = new Uint8ClampedArray(buffer.length)
//   for (let i = 0; i < buffer.length; i++) {
//     data[i] = buffer[i]
//   }
//   return data
// }
